import caffe
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import cv2
caffe.set_mode_cpu()
caffe_root="./"
model_def = caffe_root + 'Deploy/deploy.prototxt'
model_weights = caffe_root + 'verify_code_iter_10000.caffemodel'
net = caffe.Net(model_def,model_weights,caffe.TEST)
binary_mean_npy = "/Users/rogerluo/Desktop/pyopnecv/caffedemo/verifycode/mean_binaryproto.npy"
mu = np.load(binary_mean_npy)
mu = mu.mean(1).mean(1)  # average over pixels to obtain the mean (BGR) pixel values
labels_file = caffe_root + 'Synset/synset_words.txt'
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2,0,1))  # move image channels to outermost dimension
transformer.set_mean('data', mu)            # subtract the dataset-mean value in each channel
transformer.set_raw_scale('data', 255)      # rescale from [0, 1] to [0, 255]
transformer.set_channel_swap('data', (2,1,0))  # swap channels from RGB to BGR
net.blobs['data'].reshape(50,3,26, 22)  # image size is 26x22

def predict_seg(image):
    transformed_image = transformer.preprocess('data', image)
    print transformed_image.shape
    net.blobs['data'].data[...] = transformed_image
    output = net.forward()
    output_prob = output['prob'][0]
    print 'predicted class is:', output_prob.argmax()
    labels = np.loadtxt(labels_file, str, delimiter='\n')
    print 'output label:', labels[output_prob.argmax()]
    return labels[output_prob.argmax()]
pic_path="/Users/rogerluo/Desktop/pyopnecv/caffedemo/verifycode/train_notsplit/1-78Us.jpg"
img = Image.open(pic_path)
plt.subplot(1,5,1);plt.imshow(img)
# image = caffe.io.load_image(pic_path)
# print image
# print cv2.imread(pic_path)
outputlabel=[]
region = [(0,0,20,26),(16,0,38,26),(34,0,56,26),(52,0,74,26)]
for i in range(0,len(region)):
    cropImg = img.crop(region[i])
    # cropImg = cropImg.resize(26,22)
    plt.subplot(1,5,i+2);plt.imshow(cropImg)
    cropImg = np.array(cropImg,dtype=np.float32)
    cropImg = cropImg/255*1.0
    print cropImg.shape
    label=predict_seg(cropImg)

    outputlabel.append(label)

print outputlabel
plt.show()

